import os
import requests

# ==== Azure API 配置 ====
AZURE_OPENAI_API_KEY = "EVlHLUCSdiDb0yVKwkm6USjrqshc2hZxyoLD5EI7PBzl3coU7ApTJQQJ99BCACYeBjFXJ3w3AAABACOGVmn4"
AZURE_OPENAI_ENDPOINT = "https://eastus-0303.openai.azure.com/"
API_VERSION = "2024-08-01-preview"

# 列出所有部署
url = f"{AZURE_OPENAI_ENDPOINT}openai/deployments?api-version={API_VERSION}"
headers = {
    "api-key": AZURE_OPENAI_API_KEY
}

response = requests.get(url, headers=headers)

if response.status_code == 200:
    deployments = response.json()
    print("=== 当前部署列表 ===")
    for dep in deployments.get("data", []):
        print(f"部署名: {dep.get('id')}, 模型: {dep.get('model')}")
else:
    print(f"请求失败: {response.status_code}, {response.text}")

from langchain_community.document_loaders import DirectoryLoader
from langchain_openai import AzureOpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import CharacterTextSplitter
import os

# ==== Azure API 配置 ====
os.environ["AZURE_OPENAI_API_KEY"] = "EVlHLUCSdiDb0yVKwkm6USjrqshc2hZxyoLD5EI7PBzl3coU7ApTJQQJ99BCACYeBjFXJ3w3AAABACOGVmn4"
os.environ["AZURE_OPENAI_ENDPOINT"] = "https://eastus-0303.openai.azure.com/"
os.environ["AZURE_OPENAI_API_VERSION"] = "2024-08-01-preview"

# 1. 加载本地文档
loader = DirectoryLoader("company_data", glob="xinxi/ziliao.txt")
documents = loader.load()

# 2. 切分文档
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
docs = text_splitter.split_documents(documents)

# 3. 创建 Azure OpenAI Embeddings
embeddings = AzureOpenAIEmbeddings(
    deployment="my-embedding-model",        # 你在 Azure 部署的 embedding 模型名
    model="text-embedding-3-small",         # 模型类型
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version=os.environ["AZURE_OPENAI_API_VERSION"]
)

# 4. 存储到 Chroma 数据库
db = Chroma.from_documents(docs, embeddings, persist_directory="company_db")
db.persist()

print("✅ 向量数据库已创建并保存到 company_db")
print("AZURE_OPENAI_API_KEY =", os.getenv("AZURE_OPENAI_API_KEY"))
print("AZURE_OPENAI_ENDPOINT =", os.getenv("AZURE_OPENAI_ENDPOINT"))

db = Chroma.from_documents(docs, embeddings, persist_directory="company_db")
db.persist()

print("✅ 向量数据库已创建并保存到 company_db")

from langchain_community.vectorstores import Chroma
from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
from langchain.prompts import PromptTemplate
import os

# ==== Azure API 配置 ====
os.environ["AZURE_OPENAI_API_KEY"] = "EVlHLUCSdiDb0yVKwkm6USjrqshc2hZxyoLD5EI7PBzl3coU7ApTJQQJ99BCACYeBjFXJ3w3AAABACOGVmn4"
os.environ["AZURE_OPENAI_ENDPOINT"] = "https://eastus-0303.openai.azure.com/"
os.environ["AZURE_OPENAI_API_VERSION"] = "2024-08-01-preview"

# ===== 1. 读取向量数据库 =====
embeddings = AzureOpenAIEmbeddings(
    deployment="text-embedding-3-large-1",     # 你部署的 embedding 模型
    model="text-embedding-3-large",
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version="2023-05-15"
)

db = Chroma(persist_directory="company_db", embedding_function=embeddings)

# ===== 2. 初始化 Azure GPT-4o-1 =====
llm = AzureChatOpenAI(
    deployment_name="gpt-4o-1",      # 你 Azure 部署的对话模型名
    model="gpt-4o",
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version=os.environ["AZURE_OPENAI_API_VERSION"],
    temperature=0
)

# ===== 3. Prompt 模板 =====
template = """你是一个智能助手。基于以下检索到的资料回答用户问题。
如果资料中没有答案，请直接说“我不知道”，不要编造。
资料内容：
{context}

用户问题：
{question}

回答：
"""
prompt = PromptTemplate(
    input_variables=["context", "question"],
    template=template
)

# ===== 4. 运行问答 =====
while True:
    question = input("\n请输入你的问题（输入 'exit' 退出）：")
    if question.lower() == "exit":
        break

    # 从数据库检索
    docs = db.similarity_search(question, k=3)
    context = "\n\n".join([doc.page_content for doc in docs])

    # 构造 prompt
    final_prompt = prompt.format(context=context, question=question)

    # 调用 GPT
    response = llm.invoke(final_prompt)
    print("\n💡 回答：", response.content)
